#pragma once

#include <iostream>
#include <string>
#include <memory>
#include <exception>
#include <cstring>
#include <cerrno>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

namespace inet
{
struct api
{
    enum 
    {
        udp = SOCK_DGRAM,
        tcp = SOCK_STREAM,
    };

    static int Socket(int proto, bool reuse_addr = true)
    {
        int fd = socket(AF_INET, proto, 0);
        if (fd < 0) 
            throw std::runtime_error(std::string("socket failed: ") + strerror(errno));

        int opt = 1;
        if (reuse_addr && setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) 
            throw std::runtime_error(std::string("setsockopt failed: ") + strerror(errno));

        return fd;
    }

    static void Bind(int sock, const std::string& ip, uint16_t port)
    {
        struct sockaddr_in local;
        memset(&local, 0, sizeof(local));

        local.sin_family      = AF_INET;
        local.sin_addr.s_addr = inet_addr(ip.c_str());
        local.sin_port        = htons(port);

        if (bind(sock, (struct sockaddr*)&local, sizeof(local)) < 0)
            throw std::runtime_error(std::string("bind error: ") + strerror(errno));
    }

    static void Bind(int sock, uint16_t port)
    {
        struct sockaddr_in local;
        memset(&local, 0, sizeof(local));

        local.sin_family      = AF_INET;
        local.sin_addr.s_addr = INADDR_ANY;
        local.sin_port        = htons(port);

        if (bind(sock, (struct sockaddr*)&local, sizeof(local)) < 0)
            throw std::runtime_error(std::string("bind error: ") + strerror(errno));
    }

    static void Listen(int sock, int backlog = 12)
    {
        if (listen(sock, backlog) < 0)
            throw std::runtime_error(std::string("listen error: ") + strerror(errno));
    }

    static void Connect(int sock, const std::string& ip, uint16_t port, int trytime = 1)
    {
        struct sockaddr_in peer;
        memset(&peer, 0, sizeof(peer));

        peer.sin_family      = AF_INET;
        peer.sin_addr.s_addr = inet_addr(ip.c_str());
        peer.sin_port        = htons(port);

        while ((trytime-- > 0) && (connect(sock, (struct sockaddr*)&peer, sizeof(peer)) < 0))
        {
            if (trytime == 0) throw std::runtime_error(std::string("connect error: ") + strerror(errno));
            sleep(1);
        }
    }

    static int Accept(int sock, std::string* ip = nullptr, uint16_t* port = nullptr)
    {
        struct sockaddr_in peer;
        socklen_t len = sizeof(peer);
        memset(&peer, 0, len);

        int fd = accept(sock, (struct sockaddr*)&peer, &len);

        if (ip)   *ip   = inet_ntoa(peer.sin_addr);
        if (port) *port = ntohs(peer.sin_port);

        if (fd < 0)
            throw std::runtime_error(std::string("accept error: ") + strerror(errno));

        return fd;
    }

    static ssize_t Recv(int sock, std::string* msg, size_t maxlen)
    {
        msg->clear();
        std::unique_ptr<char[]> buf(new char[maxlen + 1]{0});

        ssize_t s = recv(sock, buf.get(), maxlen, 0);
        if (s > 0)
        {
            buf[s] = 0;
            *msg = buf.get();
        }
        return s;
    }

    static ssize_t Recv(int sock, char* buf, size_t maxlen)
    {
        ssize_t s = recv(sock, buf, maxlen, 0);
        if (s > 0)
        {
            if (s < maxlen)
            {
                buf[s] = 0;
            }
        }
        return s;
    }

    static ssize_t Send(int sock, const std::string& msg)
    {
        ssize_t total = 0;

        while (total < msg.size())
        {
            ssize_t s = send(sock, msg.c_str() + total, msg.size() - total, 0);

            if (s > 0) total += s;
            else return s;
        }
        return total;
    }

    static ssize_t Recvfrom(int sock, std::string* msg, size_t maxlen, std::string* ip = nullptr, uint16_t* port = nullptr)
    {
        msg->clear();
        std::unique_ptr<char[]> buf(new char[maxlen]{0});

        struct sockaddr_in peer;
        socklen_t len = sizeof(peer);
        memset(&peer, 0, len);
        ssize_t s = recvfrom(sock, buf.get(), maxlen, 0, (struct sockaddr*)&peer, &len);

        if (s > 0) 
        {
            buf[s] = 0;
            *msg = buf.get();
        }
        if (ip)   *ip   = inet_ntoa(peer.sin_addr);
        if (port) *port = ntohs(peer.sin_port);

        return s;
    }

    static ssize_t Recvfrom(int sock, char* buf, size_t maxlen, std::string* ip = nullptr, uint16_t* port = nullptr)
    {
        struct sockaddr_in peer;
        socklen_t len = sizeof(peer);
        memset(&peer, 0, len);
        ssize_t s = recvfrom(sock, buf, maxlen, 0, (struct sockaddr*)&peer, &len);

        if (s > 0) 
        {
            buf[s] = 0;
        }
        if (ip)   *ip   = inet_ntoa(peer.sin_addr);
        if (port) *port = ntohs(peer.sin_port);

        return s;
    }

    static ssize_t Sendto(int sock, const std::string& msg, const std::string& ip, uint16_t port)
    {
        struct sockaddr_in peer;
        memset(&peer, 0, sizeof(peer));

        peer.sin_family      = AF_INET;
        peer.sin_addr.s_addr = inet_addr(ip.c_str());
        peer.sin_port        = htons(port);

        ssize_t s = sendto(sock, msg.c_str(), msg.size(), 0, (struct sockaddr*)&peer, sizeof(peer));
        return s;
    }
};

namespace tcp
{
class server
{
public:
    server(uint16_t port, int backlog = 12)
        : _sock(-1), _port(port), _backlog(backlog)
    {
        __init__();
    }

    server(const std::string& ip, uint16_t port, int backlog = 12)
        : _sock(-1), _ip(ip), _port(port), _backlog(backlog)
    {
        __init__();
    }

    int accept(std::string* cip = nullptr, uint16_t* cport = nullptr) { return inet::api::Accept(_sock, cip, cport); }

    ssize_t recv(int sock, std::string* msg, size_t maxlen) { return inet::api::Recv(sock, msg, maxlen); }
    ssize_t recv(int sock,        char* buf, size_t maxlen) { return inet::api::Recv(sock, buf, maxlen); }

    ssize_t send(int sock, const std::string& msg) { return inet::api::Send(sock, msg); }

    ~server() { close(_sock); }

public:
    int         sock()    const { return _sock;    }
    std::string ip()      const { return _ip;      }
    uint16_t    port()    const { return _port;    }
    int         backlog() const { return _backlog; }

private:
    void __init__()
    {
        _sock = inet::api::Socket(inet::api::tcp);

        if (_ip.empty()) inet::api::Bind(_sock, _port);
        else             inet::api::Bind(_sock, _ip, _port);

        inet::api::Listen(_sock, _backlog);
    }

protected:
    int _sock;
    std::string _ip;
    uint16_t _port;
    int _backlog;
};

class client
{
public:
    client(const std::string& svr_ip, uint16_t svr_port, int trytime = 1)
        : _sock(-1), _sip(svr_ip), _sport(svr_port), _trytime(trytime)
    {
        _sock = inet::api::Socket(inet::api::tcp);
        inet::api::Connect(_sock, _sip, _sport, _trytime);
    }

    ssize_t send(int sock, const std::string& msg) { return inet::api::Send( sock, msg); }
    ssize_t send(          const std::string& msg) { return inet::api::Send(_sock, msg); }
    ssize_t recv(int sock, std::string* msg, size_t maxlen) { return inet::api::Recv( sock, msg, maxlen); }
    ssize_t recv(          std::string* msg, size_t maxlen) { return inet::api::Recv(_sock, msg, maxlen); }
    ssize_t recv(int sock,        char* buf, size_t maxlen) { return inet::api::Recv( sock, buf, maxlen); }
    ssize_t recv(                 char* buf, size_t maxlen) { return inet::api::Recv(_sock, buf, maxlen); }

    ~client() { close(_sock); }

public:
    int         sock()  const { return _sock;  }
    std::string sip()   const { return _sip;   }
    uint16_t    sport() const { return _sport; }

protected:
    int _sock;
    std::string _sip;
    uint16_t _sport;
    int _trytime;
};
}

namespace udp
{
class server
{
public:
    server(uint16_t port) : _sock(-1), _port(port)
    {
        __init__();
    }

    server(const std::string& ip, uint16_t port) : _sock(-1), _ip(ip), _port(port)
    {
        __init__();
    }

    ssize_t sendto(const std::string& msg, const std::string& cip, uint16_t cport)
    {
        return inet::api::Sendto(_sock, msg, cip, cport);
    }

    ssize_t recvfrom(std::string* msg, size_t maxlen, std::string* cip = nullptr, uint16_t* cport = nullptr)
    {
        return inet::api::Recvfrom(_sock, msg, maxlen, cip, cport);
    }
    ssize_t recvfrom(       char* buf, size_t maxlen, std::string* cip = nullptr, uint16_t* cport = nullptr)
    {
        return inet::api::Recvfrom(_sock, buf, maxlen, cip, cport);
    }

    ~server() { close(_sock); }

public:
    int         sock()    const { return _sock;    }
    std::string ip()      const { return _ip;      }
    uint16_t    port()    const { return _port;    }

private:
    void __init__()
    {
        _sock = inet::api::Socket(inet::api::udp);

        if (_ip.empty()) inet::api::Bind(_sock,      _port);
        else             inet::api::Bind(_sock, _ip, _port);
    }

protected:
    int _sock;
    std::string _ip;
    uint16_t _port;
};

class client
{
public:
    client(const std::string& svr_ip, uint16_t svr_port)
        : _sock(-1), _sip(svr_ip), _sport(svr_port)
    {
        _sock = inet::api::Socket(inet::api::udp);
    }

    ssize_t sendto(const std::string& msg) { return inet::api::Sendto(_sock, msg, _sip, _sport); }

    ssize_t recvfrom(std::string* msg, size_t maxlen) { return inet::api::Recvfrom(_sock, msg, maxlen); }
    ssize_t recvfrom(       char* buf, size_t maxlen) { return inet::api::Recvfrom(_sock, buf, maxlen); }

    ~client() { close(_sock); }

public:
    int         sock()  const { return _sock;  }
    std::string sip()   const { return _sip;   }
    uint16_t    sport() const { return _sport; }

protected:
    int _sock;
    std::string _sip;
    uint16_t _sport;
};
}

}
